{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "from miasm.analysis.machine import Machine\n",
    "from miasm.arch.x86.arch import mn_x86\n",
    "from miasm.ir.symbexec import SymbolicExecutionEngine\n",
    "from miasm.expression.expression import ExprCond, ExprId, ExprInt, ExprMem \n",
    "from miasm.expression.simplifications import expr_simp\n",
    "from miasm.arch.x86.regs import *\n",
    "from miasm.core import parse_asm, asmblock\n",
    "from miasm.core.locationdb import LocationDB\n",
    "from miasm.analysis.binary import Container\n",
    "from future.utils import viewitems\n",
    "from miasm.loader.strpatchwork import *\n",
    "from miasm.ir.translators.translator import Translator\n",
    "import warnings\n",
    "import z3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_path_feasibility(conds):\n",
    "    solver = z3.Solver()\n",
    "    for lval, rval in conds:\n",
    "        z3_cond = Translator.to_language(\"z3\").from_expr(lval)\n",
    "        solver.add(z3_cond == int(rval.arg))\n",
    "\n",
    "    rslt = solver.check()\n",
    "\n",
    "    if rslt == z3.sat:\n",
    "        return True\n",
    "    else:\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FinalState:\n",
    "    def __init__(self, result, sym, path_conds, path_history):\n",
    "        self.result = result\n",
    "        self.sb = sym\n",
    "        self.path_conds = path_conds\n",
    "        self.path_history = path_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def explore(ir, start_addr, start_symbols, \n",
    "        ircfg, cond_limit=30, uncond_limit=100, \n",
    "        lbl_stop=None, final_states=[]):\n",
    "\n",
    "    def codepath_walk(addr, symbols, conds, depth, final_states, path):\n",
    "\n",
    "        if depth >= cond_limit:\n",
    "            warnings.warn(\"'depth' is over the cond_limit :%d\"%(depth))\n",
    "            return \n",
    "\n",
    "        sb = SymbolicExecutionEngine(ir, symbols)\n",
    "\n",
    "        for _ in range(uncond_limit):\n",
    "\n",
    "            if isinstance(addr, ExprInt): \n",
    "                if addr == lbl_stop:\n",
    "                    final_states.append(FinalState(True, sb, conds, path))\n",
    "                    return\n",
    "\n",
    "            path.append(addr)\n",
    "\n",
    "            pc = sb.run_block_at(ircfg, addr)\n",
    "\n",
    "            if isinstance(pc, ExprCond): \n",
    "    \n",
    "                # Calc the condition to take true or false paths\n",
    "                cond_true  = {pc.cond: ExprInt(1, 32)}\n",
    "                cond_false = {pc.cond: ExprInt(0, 32)}\n",
    "\n",
    "                # The destination addr of the true or false paths\n",
    "                addr_true  = expr_simp(\n",
    "                    sb.eval_expr(pc.replace_expr(cond_true), {}))\n",
    "\n",
    "                addr_false = expr_simp(\n",
    "                    sb.eval_expr(pc.replace_expr(cond_false), {}))\n",
    "\n",
    "                # Need to add the path conditions to reach this point\n",
    "                conds_true = list(conds) + list(cond_true.items())\n",
    "                conds_false = list(conds) + list(cond_false.items())\n",
    "\n",
    "                # TODO:\n",
    "                # Please add some code to complete the SMT-based path explore\n",
    "                # Hint1: use check_path_feasibility()\n",
    "                # Hint2: Do not forget to add the current state to the final_states \n",
    "                # when the path condition is infeasible. \n",
    "                # e.g., final_states.append(FinalState(False, sb, conds_true/conds_false, path))\n",
    "                #\n",
    "                # From here --------------------------------\n",
    "                \n",
    "                if check_path_feasibility(conds_true):\n",
    "                    codepath_walk(\n",
    "                        addr_true, sb.symbols.copy(), \n",
    "                        conds_true, depth + 1, final_states, list(path))\n",
    "                else:\n",
    "                    final_states.append(FinalState(False, sb, conds_true, path))\n",
    "                    \n",
    "                if check_path_feasibility(conds_false):\n",
    "                    codepath_walk(\n",
    "                        addr_false, sb.symbols.copy(), \n",
    "                        conds_false, depth + 1, final_states, list(path))\n",
    "                else:\n",
    "                    final_states.append(FinalState(False, sb, conds_false, path))\n",
    "\n",
    "                # To here --------------------------------\n",
    "\n",
    "                return\n",
    "            else:\n",
    "                addr = expr_simp(sb.eval_expr(pc))\n",
    "\n",
    "        final_states.append(FinalState(True, sb, conds, path))\n",
    "        return \n",
    "\n",
    "    return codepath_walk(start_addr, start_symbols, [], 0, final_states, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble code\n",
    "loc_db = LocationDB()\n",
    "asmcfg = parse_asm.parse_txt(mn_x86, 32, ''' \n",
    "main:\n",
    "    PUSH EBP\n",
    "    MOV EBP, ESP\n",
    "    MOV ECX, 0x23\n",
    "    MOV EDX, EAX\n",
    "    MUL EDX\n",
    "    CMP EAX, -1\n",
    "    JNZ label\n",
    "    MOV DWORD PTR [0xDEADBEEF], ECX\n",
    "\n",
    "label:\n",
    "    MOV ECX, 0x4\n",
    "    MOV EAX, ECX\n",
    "    POP EBP\n",
    "    RET\n",
    "''', loc_db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc_db.set_location_offset(loc_db.get_name_location('main'), 0x0)\n",
    "\n",
    "patches = asmblock.asm_resolve_final(mn_x86, asmcfg)\n",
    "patch_worker = StrPatchwork()\n",
    "for offset, raw in patches.items():\n",
    "    patch_worker[offset] = raw\n",
    "    print('%08x'%(offset), mn_x86.dis(raw, 32))\n",
    "\n",
    "cont = Container.from_string(array_tobytes(patch_worker.s), loc_db)\n",
    "\n",
    "machine = Machine('x86_32')\n",
    "mdis = machine.dis_engine(cont.bin_stream, loc_db=loc_db)\n",
    "\n",
    "asmcfg2 = mdis.dis_multiblock(0)\n",
    "\n",
    "ir_arch = machine.ira(loc_db)\n",
    "ircfg = ir_arch.new_ircfg_from_asmcfg(asmcfg2)\n",
    "\n",
    "for lbl, irb in viewitems(ircfg.blocks):\n",
    "    print(irb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "symbols_init =  {\n",
    "    ExprMem(ExprId('ESP_init', 32), 32) : ExprInt(0xdeadbeef, 32)\n",
    "}\n",
    "\n",
    "for i, r in enumerate(all_regs_ids):\n",
    "    symbols_init[r] = all_regs_ids_init[i]\n",
    "\n",
    "final_states = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "explore(ir_arch, \n",
    "        0, \n",
    "        symbols_init, \n",
    "        ircfg, \n",
    "        lbl_stop=0xdeadbeef, \n",
    "        final_states=final_states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show results\n",
    "print('final states:', len(final_states))\n",
    "\n",
    "for final_state in final_states:\n",
    "    if final_state.result:\n",
    "        print('Feasible path:','->'.join([str(x) for x in final_state.path_history]))\n",
    "        print('\\t',final_state.path_conds)\n",
    "    else:\n",
    "        print('Infeasible path:','->'.join([str(x) for x in final_state.path_history]))\n",
    "        print('\\t',final_state.path_conds)\n",
    "\n",
    "    final_state.sb.dump(ids=False)\n",
    "    print('')\n",
    "    \n",
    "assert len(final_states) == 2, '# of Final States is incorrect'\n",
    "assert len([x for x in final_states if x.result == True]) == 1, '# of feasible is incorrect'\n",
    "assert len([x for x in final_states if x.result == False]) == 1, '# of infeasible is incorrect'\n",
    "\n",
    "print('Congrats!')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
